package com.itcam.avalbackend.config;

import com.itcam.avalbackend.rag.CosAppDocumentLoader;
import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import java.util.List;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

@Configuration
public class PgVectorVectorStoreConfig {

    @Resource
    private CosAppDocumentLoader cosAppDocumentLoader;

    @Bean
    public VectorStore pgVectorVectorStore(JdbcTemplate jdbcTemplate, @Qualifier("dashscopeEmbeddingModel") EmbeddingModel dashscopeEmbeddingModel) {
        // 使用 PgVectorStore 的构建器模式创建 VectorStore 实例
        VectorStore vectorStore = PgVectorStore.builder(jdbcTemplate, dashscopeEmbeddingModel)
                // 向量的维度，可选参数，默认值为模型维度或 1536
                .dimensions(1536)
                // 向量距离计算类型，可选参数，默认使用余弦距离
                .distanceType(COSINE_DISTANCE)
                // 索引类型，可选参数，默认使用 HNSW（Hierarchical Navigable Small World）索引
                .indexType(HNSW)
                // 是否初始化数据库模式，可选参数，默认不初始化
                .initializeSchema(true)
                // 数据库模式名称，可选参数，默认使用 "public" 模式
                .schemaName("public")
                // 存储向量的表名，可选参数，默认表名为 "vector_store"
                .vectorTableName("vector_store")
                // 最大文档批量处理数量，可选参数，默认值为 10000
                .maxDocumentBatchSize(10000)
                .build();
        // 加载文档
        List<Document> documents = cosAppDocumentLoader.loadMarkdowns();
        vectorStore.add(documents);
        return vectorStore;
    }
}